package com.feiyue.learning.advance.nio;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * TCP Server 服务端
 * @author  feiyue
 * @date  2020/4/18
 */
public class NioReceiveServer {
    private Charset charset = Charset.forName("UTF-8");
    private ByteBuffer buffer = ByteBuffer.allocate(9);
    private static final String DEST_DIRECTORY = "C:\\Users\\zhou.xm\\Desktop";
    Map<SelectableChannel, Client> clientMap = new HashMap<>(8);

    public static void main(String[] args) throws IOException {
        NioReceiveServer server = new NioReceiveServer();
        server.startServer();
    }

    private void startServer() throws IOException {
        // 获取选择器
        Selector selector = Selector.open();

        // 获取通道
        ServerSocketChannel serverChannel = ServerSocketChannel.open();
        ServerSocket serverSocket = serverChannel.socket();
        serverChannel.configureBlocking(false);
        serverSocket.bind(new InetSocketAddress(80));

        // 将通道注册到选择器并注册 IO 事件为 `接收新连接`
        serverChannel.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("serverChannel is listening...");

        // 选择感兴趣的 IO 就绪事件
        while (selector.select() > 0) {
            Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();
            while (keyIterator.hasNext()) {
                SelectionKey key = keyIterator.next();
                if (key.isAcceptable()) {
                    // 获取选择键中注册的服务器通道
                    ServerSocketChannel server = (ServerSocketChannel) key.channel();

                    // 获取新客户端连接，并注册到选择器中以便后续读取该通道数据
                    SocketChannel socketChannel = server.accept();
                    if (socketChannel == null) {
                        continue;
                    }
                    socketChannel.configureBlocking(false);
                    SelectionKey selectionKey = socketChannel.register(selector, SelectionKey.OP_READ);

                    // 为每个客户端连接生成一个 client 对象, 放到 map, 以便后续使用
                    Client client = new Client();
                    client.remoteAddress = (InetSocketAddress) socketChannel.getRemoteAddress();
                    clientMap.put(socketChannel, client);
                    System.out.println(socketChannel.getRemoteAddress() + "连接成功...");
                } else if (key.isReadable()) {
                    processData(key);
                }

                // 已选择的键不会被删除, 为了避免下一次再次被选中, 手动删除
                keyIterator.remove();
            }
        }
    }

    private void processData(SelectionKey key) throws IOException {
        Client client = clientMap.get(key.channel());
        SocketChannel socketChannel = (SocketChannel)key.channel();

        int num = 0;
        try {
            buffer.clear();
            while ((num = socketChannel.read(buffer)) > 0) {
                buffer.flip();
                if (client.fileName == null) {
                    String fileName = charset.decode(buffer).toString();
                    File directory = new File(DEST_DIRECTORY);
                    if (! directory.exists()) {
                        directory.mkdir();
                    }
                    client.fileName = fileName;
                    String fullName = directory.getAbsolutePath() + File.separator + fileName;
                    System.out.println("NIO 传输目标文件：" + fullName);

                    File file = new File(fullName);
                    FileChannel fileChannel = new FileOutputStream(file).getChannel();
                    client.outChannel = fileChannel;
                } else if (client.fileLength == 0) {
                    long fileLength = buffer.getLong();
                    client.fileLength = fileLength;
                    client.startTime = System.currentTimeMillis();
                    System.out.println("NIO 开始传输文件...");
                } else {
                    // 客户端发送过来的是文件内容，写入文件
                    client.outChannel.write(buffer);
                }
                buffer.clear();
            }
            key.cancel();
        } catch (Exception e) {
            key.cancel();
            e.printStackTrace();
            return;
        }

        // 读取数量 -1, 表示传输结束
        if (num == -1) {
            client.outChannel.close();
            System.out.println("传输完毕");
            key.cancel();
            System.out.println("文件接收成功, fileName: " + client.fileName + ", size: " + client.fileLength);
            System.out.println("文件传输耗时：" + (System.currentTimeMillis() - client.startTime));
        }
    }

    // 内部类，服务器端保存客户端对象
    static class Client {
        String fileName;
        long fileLength;
        long startTime;
        InetSocketAddress remoteAddress;
        FileChannel outChannel;
    }
}
